package com.LeeCode;

import java.util.Deque;
import java.util.LinkedList;

/**
 * 统计全 1 子矩形
 */

public class Code1504 {
    public static void main(String[] args) {

    }

    public int numSubmat(int[][] mat) {
        int m = mat.length;
        int n = mat[0].length;
        int result = 0;

        // heights[i][j] 表示位置 (i,j) 向上连续1的高度
        int[][] heights = new int[m][n];

        // 计算每个位置的向上高度
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                if (mat[i][j] == 1) {
                    heights[i][j] = (i == 0) ? 1 : heights[i - 1][j] + 1;
                }
            }
        }

        // 对每一行计算子矩形数量
        for (int i = 0; i < m; i++) {
            result += countSubmatricesInRow(heights[i]);
        }

        return result;
    }

    private int countSubmatricesInRow(int[] heights) {
        int n = heights.length;
        Deque<Integer> stack = new LinkedList<>();
        int[] dp = new int[n];  // dp[i] 表示以第i列为右边界的子矩形数量
        int result = 0;

        for (int i = 0; i < n; i++) {
            while (!stack.isEmpty() && heights[stack.peek()] >= heights[i]) {
                stack.pop();
            }

            if (stack.isEmpty()) {
                dp[i] = heights[i] * (i + 1);
            } else {
                int prevIndex = stack.peek();
                dp[i] = dp[prevIndex] + heights[i] * (i - prevIndex);
            }

            result += dp[i];
            stack.push(i);
        }

        return result;
    }
}
